import torch
from PIL import Image
import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import os
from threading import Thread
MODEL_NAME = "hf-models/internlm2-chat-7b"
TITLE = "<h1><center>internlm2-chat-7b</center></h1>"
DESCRIPTION = ''
CSS = """
.duplicate-button {
  margin: auto !important;
  color: white !important;
  background: black !important;
  border-radius: 100vh !important;
}
"""
model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16,
        trust_remote_code=True,
        ).to(0).eval()
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME,trust_remote_code=True)
def stream_chat(message, history):
    messages = [{
              "role": "system",
              "content":
              "你是一个全能的助手，使用中文, 不要附带 JSON 格式回答问题"
            }]
    
    for msg in history:
        messages.append({"role": "user", "content": str(msg[0])})
        messages.append({"role": "assistant", "content": str(msg[1])})
        
    messages.append({"role": "user", "content": str (message)})
    res = model.stream_chat(tokenizer, message)
    length = 0
    for response, history in res:
        # 逐行输出
        # yield response[length:]
        # length=len(response)
        yield response
chatbot = gr.Chatbot(height=450)
with gr.Blocks(css=CSS, theme=gr.themes.Glass()) as demo:
    gr.HTML(TITLE)
    gr.HTML(DESCRIPTION)
    gr.ChatInterface(
        fn=stream_chat,
        chatbot=chatbot,
        fill_height=True,
        examples=[
            ["鲁迅和周树人什么关系"],
            ["从前有一头牛，这头牛后面有什么？"],
            ["以红楼梦的行文风格写一张委婉的请假条。不少于320字。"],
            ["人生的意义是什么"],
            ["js 实现python 的 range(10)"],
            ["我的蓝牙耳机坏了，我该去看牙科还是耳鼻喉科？"],
            ["鸡和兔子同笼，头共 10，脚共28，鸡兔子各有几只？"],
        ],
        cache_examples=False,
    )
if __name__ == "__main__":
    demo.queue("auto")
    demo.launch(debug=True)
# for response, history in model.stream_chat(tokenizer, [{
#                 "role": "system",
#                 "content":
#                 "你是一个聪明的助手"
#                 },{"role": "user", "content": "讲个笑话"}], history=[]):
#     print(response + "\n")
#     length = len(response)